# https://gitee.com/yueyinqiu5990/tj12413601/blob/master/assignment3/question2/main_guess.py
import matplotlib.pyplot
import matplotlib.widgets
import torch
import torch.utils.data

import datasets
import models


def main():
    dataset = datasets.EpsDataset()
    data_loader = torch.utils.data.DataLoader(dataset, batch_size=len(dataset))
    data_x, data_y = next(iter(data_loader))

    matplotlib.pyplot.scatter(data_x, data_y[:, 0], c="moccasin")
    matplotlib.pyplot.scatter(data_x, data_y[:, 1], c="thistle")
    matplotlib.pyplot.ylim(torch.min(data_y), torch.max(data_y))

    x = torch.linspace(0, 90, 1000)
    line1 = matplotlib.pyplot.plot(x, x, c="orange")[0]
    line2 = matplotlib.pyplot.plot(x, x, c="purple")[0]

    matplotlib.pyplot.subplots_adjust(bottom=0.30)
    slider_epsilon_infinity = matplotlib.widgets.Slider(
        matplotlib.pyplot.axes((0.2, 0.20, 0.65, 0.03)),
        "epsilon_infinity",
        valmin=-1, valmax=3, valinit=0.962)
    slider_delta_epsilon = matplotlib.widgets.Slider(
        matplotlib.pyplot.axes((0.2, 0.15, 0.65, 0.03)),
        "delta_epsilon",
        valmin=3, valmax=7, valinit=5.34)
    slider_tau = matplotlib.widgets.Slider(
        matplotlib.pyplot.axes((0.2, 0.10, 0.65, 0.03)),
        "tau",
        valmin=-2, valmax=2, valinit=-0.083)

    def update(_=None):
        model = models.DebyeModel(
            slider_epsilon_infinity.val,
            slider_delta_epsilon.val,
            slider_tau.val)
        y = torch.detach(model(x))
        line1.set_ydata(y[:, 0])
        line2.set_ydata(y[:, 1])

    slider_epsilon_infinity.on_changed(update)
    slider_delta_epsilon.on_changed(update)
    slider_tau.on_changed(update)
    update()

    matplotlib.pyplot.show()


if __name__ == "__main__":
    main()
